/*
Copyright 2020 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package clusterapi

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	v1 "k8s.io/api/core/v1"
	resourceapi "k8s.io/api/resource/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/client-go/tools/cache"
	"k8s.io/utils/ptr"
)

const (
	cpuStatusKey             = "cpu"
	memoryStatusKey          = "memory"
	nvidiaGpuStatusKey       = "nvidia.com/gpu"
	architectureStatusKey    = "architecture"
	operatingSystemStatusKey = "operatingSystem"

	arm64 = "arm64"
	linux = "linux"
)

func TestSetSize(t *testing.T) {
	initialReplicas := 1
	updatedReplicas := 5
	finalReplicas := 0

	test := func(t *testing.T, testConfig *TestConfig) {
		controller := NewTestMachineController(t)
		defer controller.Stop()
		controller.AddTestConfigs(testConfig)

		testResource := testConfig.machineSet
		if testConfig.machineDeployment != nil {
			testResource = testConfig.machineDeployment
		}

		sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
		if err != nil {
			t.Fatal(err)
		}

		gvr, err := sr.GroupVersionResource()
		if err != nil {
			t.Fatal(err)
		}

		// First update to updatedReplicas
		err = sr.SetSize(updatedReplicas)
		if err != nil {
			t.Fatal(err)
		}

		s, err := sr.controller.managementScaleClient.Scales(testResource.GetNamespace()).
			Get(context.TODO(), gvr.GroupResource(), testResource.GetName(), metav1.GetOptions{})
		if err != nil {
			t.Fatalf("error getting scale subresource: %v", err)
		}

		if s.Spec.Replicas != int32(updatedReplicas) {
			t.Errorf("expected %v, got: %v", updatedReplicas, s.Spec.Replicas)
		}

		replicas, found, err := unstructured.NestedInt64(sr.unstructured.Object, "spec", "replicas")
		if err != nil {
			t.Fatal(err)
		}
		if !found {
			t.Fatal("replicas = 0")
		}
		if replicas != int64(updatedReplicas) {
			t.Errorf("expected %v, got: %v", updatedReplicas, replicas)
		}

		// Second update to finalReplicas
		err = sr.SetSize(finalReplicas)
		if err != nil {
			t.Fatal(err)
		}

		s, err = sr.controller.managementScaleClient.Scales(testResource.GetNamespace()).
			Get(context.TODO(), gvr.GroupResource(), testResource.GetName(), metav1.GetOptions{})
		if err != nil {
			t.Fatalf("error getting scale subresource: %v", err)
		}

		if s.Spec.Replicas != int32(finalReplicas) {
			t.Errorf("expected %v, got: %v", finalReplicas, s.Spec.Replicas)
		}

		replicas, found, err = unstructured.NestedInt64(sr.unstructured.Object, "spec", "replicas")
		if err != nil {
			t.Fatal(err)
		}
		if !found {
			t.Fatal("replicas = 0")
		}
		if replicas != int64(finalReplicas) {
			t.Errorf("expected %v, got: %v", finalReplicas, replicas)
		}
	}

	annotations := map[string]string{
		nodeGroupMinSizeAnnotationKey: "0",
		nodeGroupMaxSizeAnnotationKey: "10",
	}
	t.Run("MachineSet", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineSet().
			WithNodeCount(initialReplicas).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig)
	})

	t.Run("MachineDeployment", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineDeployment().
			WithNodeCount(initialReplicas).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig)
	})
}

func TestReplicas(t *testing.T) {
	initialReplicas := 1
	updatedReplicas := 5

	test := func(t *testing.T, testConfig *TestConfig) {
		controller := NewTestMachineController(t)
		defer controller.Stop()
		controller.AddTestConfigs(testConfig)

		testResource := testConfig.machineSet
		if testConfig.machineDeployment != nil {
			testResource = testConfig.machineDeployment
		}

		sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
		if err != nil {
			t.Fatal(err)
		}

		gvr, err := sr.GroupVersionResource()
		if err != nil {
			t.Fatal(err)
		}

		i, err := sr.Replicas()
		if err != nil {
			t.Fatal(err)
		}

		if i != initialReplicas {
			t.Errorf("expected %v, got: %v", initialReplicas, i)
		}

		// fetch and update machineSet
		s, err := sr.controller.managementScaleClient.Scales(testResource.GetNamespace()).
			Get(context.TODO(), gvr.GroupResource(), testResource.GetName(), metav1.GetOptions{})
		if err != nil {
			t.Fatal(err)
		}

		s.Spec.Replicas = int32(updatedReplicas)

		ch := make(chan error)
		checkDone := func(obj interface{}) (bool, error) {
			u, ok := obj.(*unstructured.Unstructured)
			if !ok {
				return false, nil
			}
			sr, err := newUnstructuredScalableResource(controller.machineController, u)
			if err != nil {
				return true, err
			}
			i, err := sr.Replicas()
			if err != nil {
				return true, err
			}
			if i != updatedReplicas {
				return true, fmt.Errorf("expected %v, got: %v", updatedReplicas, i)
			}
			return true, nil
		}
		handler := cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				match, err := checkDone(obj)
				if match {
					ch <- err
				}
			},
			UpdateFunc: func(oldObj, newObj interface{}) {
				match, err := checkDone(newObj)
				if match {
					ch <- err
				}
			},
		}

		if _, err := controller.machineSetInformer.Informer().AddEventHandler(handler); err != nil {
			t.Fatal(err)
		}
		if _, err := controller.machineDeploymentInformer.Informer().AddEventHandler(handler); err != nil {
			t.Fatal(err)
		}

		_, err = sr.controller.managementScaleClient.Scales(testResource.GetNamespace()).
			Update(context.TODO(), gvr.GroupResource(), s, metav1.UpdateOptions{})
		if err != nil {
			t.Fatal(err)
		}

		lastErr := fmt.Errorf("no updates received yet")
		for lastErr != nil {
			select {
			case err = <-ch:
				lastErr = err
			case <-time.After(1 * time.Second):
				t.Fatal(fmt.Errorf("timeout while waiting for update. Last error was: %v", lastErr))
			}
		}
	}

	t.Run("MachineSet", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineSet().
			WithNodeCount(initialReplicas).
			Build()
		test(t, testConfig)
	})

	t.Run("MachineDeployment", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineDeployment().
			WithNodeCount(initialReplicas).
			Build()
		test(t, testConfig)
	})
}

func TestSetSizeAndReplicas(t *testing.T) {
	initialReplicas := 1
	updatedReplicas := 5

	test := func(t *testing.T, testConfig *TestConfig) {
		controller := NewTestMachineController(t)
		defer controller.Stop()
		controller.AddTestConfigs(testConfig)

		testResource := testConfig.machineSet
		if testConfig.machineDeployment != nil {
			testResource = testConfig.machineDeployment
		}

		sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
		if err != nil {
			t.Fatal(err)
		}

		i, err := sr.Replicas()
		if err != nil {
			t.Fatal(err)
		}

		if i != initialReplicas {
			t.Errorf("expected %v, got: %v", initialReplicas, i)
		}

		err = sr.SetSize(updatedReplicas)
		if err != nil {
			t.Fatal(err)
		}

		i, err = sr.Replicas()
		if err != nil {
			t.Fatal(err)
		}

		if i != updatedReplicas {
			t.Errorf("expected %v, got: %v", updatedReplicas, i)
		}
	}

	annotations := map[string]string{
		nodeGroupMinSizeAnnotationKey: "1",
		nodeGroupMaxSizeAnnotationKey: "10",
	}

	t.Run("MachineSet", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineSet().
			WithNodeCount(initialReplicas).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig)
	})

	t.Run("MachineDeployment", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineDeployment().
			WithNodeCount(initialReplicas).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig)
	})
}

func TestAnnotations(t *testing.T) {
	cpuQuantity := resource.MustParse("2")
	memQuantity := resource.MustParse("1024Mi")
	diskQuantity := resource.MustParse("100Gi")
	gpuQuantity := resource.MustParse("1")
	maxPodsQuantity := resource.MustParse("42")
	expectedTaints := []v1.Taint{{Key: "key1", Effect: v1.TaintEffectNoSchedule, Value: "value1"}, {Key: "key2", Effect: v1.TaintEffectNoExecute, Value: "value2"}}
	testNodeName := "test-node"
	draDriver := "test-driver"
	expectedResourceSlice := &resourceapi.ResourceSlice{
		ObjectMeta: metav1.ObjectMeta{
			Name: testNodeName + "-" + draDriver,
		},
		Spec: resourceapi.ResourceSliceSpec{
			Driver:   draDriver,
			NodeName: &testNodeName,
			Pool: resourceapi.ResourcePool{
				Name: testNodeName,
			},
			Devices: []resourceapi.Device{
				{
					Name: "gpu-0",
					Attributes: map[resourceapi.QualifiedName]resourceapi.DeviceAttribute{
						"type": {
							StringValue: ptr.To(GpuDeviceType),
						},
					},
				},
			},
		},
	}
	annotations := map[string]string{
		cpuKey:          cpuQuantity.String(),
		memoryKey:       memQuantity.String(),
		diskCapacityKey: diskQuantity.String(),
		gpuCountKey:     gpuQuantity.String(),
		maxPodsKey:      maxPodsQuantity.String(),
		taintsKey:       "key1=value1:NoSchedule,key2=value2:NoExecute",
		labelsKey:       "key3=value3,key4=value4,key5=value5",
		draDriverKey:    draDriver,
	}

	test := func(t *testing.T, testConfig *TestConfig, testResource *unstructured.Unstructured) {
		controller := NewTestMachineController(t)
		defer controller.Stop()
		controller.AddTestConfigs(testConfig)

		sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
		if err != nil {
			t.Fatal(err)
		}

		if cpu, err := sr.InstanceCPUCapacityAnnotation(); err != nil {
			t.Fatal(err)
		} else if cpuQuantity.Cmp(cpu) != 0 {
			t.Errorf("expected %v, got %v", cpuQuantity, cpu)
		}

		if mem, err := sr.InstanceMemoryCapacityAnnotation(); err != nil {
			t.Fatal(err)
		} else if memQuantity.Cmp(mem) != 0 {
			t.Errorf("expected %v, got %v", memQuantity, mem)
		}

		if disk, err := sr.InstanceEphemeralDiskCapacityAnnotation(); err != nil {
			t.Fatal(err)
		} else if diskQuantity.Cmp(disk) != 0 {
			t.Errorf("expected %v, got %v", diskQuantity, disk)
		}

		if gpu, err := sr.InstanceGPUCapacityAnnotation(); err != nil {
			t.Fatal(err)
		} else if gpuQuantity.Cmp(gpu) != 0 {
			t.Errorf("expected %v, got %v", gpuQuantity, gpu)
		}

		if maxPods, err := sr.InstanceMaxPodsCapacityAnnotation(); err != nil {
			t.Fatal(err)
		} else if maxPodsQuantity.Cmp(maxPods) != 0 {
			t.Errorf("expected %v, got %v", maxPodsQuantity, maxPods)
		}

		if resourceSlices, err := sr.InstanceResourceSlices(testNodeName); err != nil {
			t.Fatal(err)
		} else {
			for _, resourceslice := range resourceSlices {
				assert.Equal(t, expectedResourceSlice, resourceslice)
			}
		}

		taints := sr.Taints()
		assert.Equal(t, expectedTaints, taints)

		labels := sr.Labels()
		assert.Len(t, labels, 3)
		assert.Equal(t, "value3", labels["key3"])
		assert.Equal(t, "value4", labels["key4"])
		assert.Equal(t, "value5", labels["key5"])
	}

	t.Run("MachineSet", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineSet().
			WithNodeCount(1).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig, testConfig.machineSet)
	})

	t.Run("MachineDeployment", func(t *testing.T) {
		testConfig := NewTestConfigBuilder().
			ForMachineDeployment().
			WithNodeCount(1).
			WithAnnotations(annotations).
			Build()
		test(t, testConfig, testConfig.machineDeployment)
	})
}

func TestCanScaleFromZero(t *testing.T) {
	testConfigs := []struct {
		name        string
		annotations map[string]string
		capacity    map[string]string
		canScale    bool
	}{
		{
			"can scale from zero",
			map[string]string{
				cpuKey:    "1",
				memoryKey: "1024Mi",
			},
			nil,
			true,
		},
		{
			"with missing CPU info cannot scale from zero",
			map[string]string{
				memoryKey: "1024Mi",
			},
			nil,
			false,
		},
		{
			"with missing Memory info cannot scale from zero",
			map[string]string{
				cpuKey: "1",
			},
			nil,
			false,
		},
		{
			"with no information cannot scale from zero",
			map[string]string{},
			nil,
			false,
		},
		{
			"with capacity in machine template can scale from zero",
			map[string]string{},
			map[string]string{
				cpuStatusKey:    "1",
				memoryStatusKey: "4G",
			},
			true,
		},
		{
			"with missing cpu capacity in machine template cannot scale from zero",
			map[string]string{},
			map[string]string{
				memoryStatusKey: "4G",
			},
			false,
		},
		{
			"with missing memory capacity in machine template cannot scale from zero",
			map[string]string{},
			map[string]string{
				cpuStatusKey: "1",
			},
			false,
		},
		{
			"with both annotations and capacity in machine template can scale from zero",
			map[string]string{
				cpuKey:    "1",
				memoryKey: "1024Mi",
			},
			map[string]string{
				cpuStatusKey:    "1",
				memoryStatusKey: "4G",
			},
			true,
		},
		{
			"with incomplete annotations and capacity in machine template cannot scale from zero",
			map[string]string{
				cpuKey: "1",
			},
			map[string]string{
				nvidiaGpuStatusKey: "1",
			},
			false,
		},
		{
			"with complete information split across annotations and capacity in machine template can scale from zero",
			map[string]string{
				cpuKey: "1",
			},
			map[string]string{
				memoryStatusKey: "4G",
			},
			true,
		},
	}

	for _, tc := range testConfigs {
		testname := fmt.Sprintf("MachineSet %s", tc.name)
		t.Run(testname, func(t *testing.T) {
			msTestConfig := NewTestConfigBuilder().
				ForMachineSet().
				WithNodeCount(1).
				WithAnnotations(tc.annotations).
				WithCapacity(tc.capacity).
				Build()
			controller := NewTestMachineController(t)
			defer controller.Stop()
			controller.AddTestConfigs(msTestConfig)

			testResource := msTestConfig.machineSet

			sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
			if err != nil {
				t.Fatal(err)
			}

			canScale := sr.CanScaleFromZero()
			if canScale != tc.canScale {
				t.Errorf("expected %v, got %v", tc.canScale, canScale)
			}
		})
	}

	for _, tc := range testConfigs {
		testname := fmt.Sprintf("MachineDeployment %s", tc.name)
		t.Run(testname, func(t *testing.T) {
			mdTestConfig := NewTestConfigBuilder().
				ForMachineDeployment().
				WithNodeCount(1).
				WithAnnotations(tc.annotations).
				WithCapacity(tc.capacity).
				Build()
			controller := NewTestMachineController(t)
			defer controller.Stop()
			controller.AddTestConfigs(mdTestConfig)

			testResource := mdTestConfig.machineDeployment

			sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
			if err != nil {
				t.Fatal(err)
			}

			canScale := sr.CanScaleFromZero()
			if canScale != tc.canScale {
				t.Errorf("expected %v, got %v", tc.canScale, canScale)
			}
		})
	}
}

func TestInstanceSystemInfo(t *testing.T) {
	// use a constant capacity as that's necessary for the business logic to consider the resource scalable
	capacity := map[string]string{
		cpuStatusKey:    "1",
		memoryStatusKey: "4G",
	}
	testConfigs := []struct {
		name         string
		nodeInfo     map[string]string
		expectedArch string
		expectedOS   string
	}{
		{
			"with no architecture or operating system in machine template's status' nodeInfo, the system info is empty",
			map[string]string{},
			"",
			"",
		},
		{
			"with architecture in machine template's status' nodeInfo, the system info is filled in the scalable resource",
			map[string]string{
				architectureStatusKey: arm64,
			},
			arm64,
			"",
		},
		{
			"with operating system in machine template's status' nodeInfo, the system info is filled in the scalable resource",
			map[string]string{
				operatingSystemStatusKey: linux,
			},
			"",
			linux,
		},
		{
			"with architecture and operating system in machine template's status' nodeInfo, the system info is filled in the scalable resource",
			map[string]string{
				architectureStatusKey:    arm64,
				operatingSystemStatusKey: linux,
			},
			arm64,
			linux,
		},
	}

	for _, tc := range testConfigs {
		testname := fmt.Sprintf("MachineSet %s", tc.name)
		t.Run(testname, func(t *testing.T) {
			mdTestConfig := NewTestConfigBuilder().
				ForMachineSet().
				WithNodeCount(1).
				WithCapacity(capacity).
				WithNodeInfo(tc.nodeInfo).
				Build()
			controller := NewTestMachineController(t)
			defer controller.Stop()
			controller.AddTestConfigs(mdTestConfig)

			testResource := mdTestConfig.machineSet

			sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
			if err != nil {
				t.Fatal(err)
			}

			sysInfo := sr.InstanceSystemInfo()
			if sysInfo.Architecture != tc.expectedArch {
				t.Errorf("expected architecture %s, got %s", tc.nodeInfo[architectureStatusKey], sysInfo.Architecture)
			}
			if sysInfo.OperatingSystem != tc.expectedOS {
				t.Errorf("expected operating system %s, got %s", tc.nodeInfo[operatingSystemStatusKey], sysInfo.OperatingSystem)
			}
		})
	}

	for _, tc := range testConfigs {
		testname := fmt.Sprintf("MachineDeployment %s", tc.name)
		t.Run(testname, func(t *testing.T) {
			mdTestConfig := NewTestConfigBuilder().
				ForMachineDeployment().
				WithNodeCount(1).
				WithCapacity(capacity).
				WithNodeInfo(tc.nodeInfo).
				Build()
			controller := NewTestMachineController(t)
			defer controller.Stop()
			controller.AddTestConfigs(mdTestConfig)

			testResource := mdTestConfig.machineDeployment

			sr, err := newUnstructuredScalableResource(controller.machineController, testResource)
			if err != nil {
				t.Fatal(err)
			}

			sysInfo := sr.InstanceSystemInfo()
			if sysInfo.Architecture != tc.expectedArch {
				t.Errorf("expected architecture %s, got %s", tc.nodeInfo[architectureStatusKey], sysInfo.Architecture)
			}

			if sysInfo.OperatingSystem != tc.expectedOS {
				t.Errorf("expected operating system %s, got %s", tc.nodeInfo[operatingSystemStatusKey], sysInfo.OperatingSystem)
			}
		})
	}
}
